/*
 * hmm_partial.cc
 *
 *  Created on: Jul 13, 2011
 *      Author: tqlong
 */

#include "hmm_partial.h"

PartialHMM::PartialHMM(int n_state, int n_observation) :
  HiddenMarkovModel(n_state, n_observation)
{

}

/** Maximum a priori state sequence (viterbi algorithm) */
void
PartialHMM::map(const seq_type &oseq, const seq_type &p_sseq, seq_type &sseq)
{
  assert(oseq.size() == p_sseq.size());

  seq_ = oseq;
  p_sseq_ = p_sseq;
  viterbi_forward_pass();
  viterbi_backward_pass();

  int T = seq_.size();
  sseq = p_sseq_;
  for (int t = 0; t < T; t++)
    if (sseq[t] == -1)
      {
        int maxi = -1;
        double maxv = LOGZERO;
        for (int i = 0; i < n_state_; i++)
          {
            double v;
            if (t == 0)
              v = logmul(logpi_[i], beta_.at(1, i));
            else if (t == T - 1)
              v = logmul(alpha_.at(T - 2, i), log(1));
            else
              v = logmul(alpha_.at(t - 1, i), beta_.at(t + 1, i));
            v = logmul(v, logb_.at(i, seq_[t]));
            if (v > maxv)
              {
                maxv = v;
                maxi = i;
              }
          }
        assert(maxi != -1);
        sseq[t] = maxi;
      }
}

void
PartialHMM::viterbi_forward_pass()
{
  int T = seq_.size();
  alpha_ = arma::zeros<mat_type>(T, n_state_);

  for (int t = 0; t < T - 1; t++)
    {
      if (p_sseq_[t] != -1) // q_t is known
        {
          for (int i = 0; i < n_state_; i++)
            alpha_.at(t, i) = loga_.at(p_sseq_[t], i);
        }
      else // q_t is unknown
        {
          for (int i = 0; i < n_state_; i++)
            {
              double maxv = LOGZERO;
              for (int q_t = 0; q_t < n_state_; q_t++)
                {
                  double v = (t == 0) ? logpi_[q_t] : alpha_.at(t - 1, q_t);
                  v = logmul(v, logb_.at(q_t, seq_[t]));
                  v = logmul(v, loga_.at(q_t, i));
                  if (v > maxv)
                    maxv = v;
                }
              alpha_.at(t, i) = maxv;
            }
        }
    }
}

void
PartialHMM::viterbi_backward_pass()
{
  int T = seq_.size();
  beta_ = arma::zeros<mat_type>(T, n_state_);
  for (int t = T - 1; t > 0; t--)
    {
      if (p_sseq_[t] != -1) // q_t is known
        {
          for (int i = 0; i < n_state_; i++)
            beta_.at(t, i) = loga_.at(i, p_sseq_[t]);
        }
      else
        {
          for (int i = 0; i < n_state_; i++)
            {
              double maxv = LOGZERO;
              for (int q_t = 0; q_t < n_state_; q_t++)
                {
                  double v = (t == T - 1) ? log(1) : beta_.at(t + 1, q_t);
                  v = logmul(v, logb_.at(q_t, seq_[t]));
                  v = logmul(v, loga_.at(i, q_t));
                  if (v > maxv)
                    maxv = v;
                }
              beta_.at(t, i) = maxv;
            }
        }
    }
}
